"""
加载和配置模型,  

"""

import subprocess
import threading



# from transformers import AutoModelForCausalLM, AutoTokenizer
# from transformers.generation import GenerationConfig

# Note: The default behavior now has injection attack prevention off.
# tokenizer = AutoTokenizer.from_pretrained("hf-models/Qwen2-72B-Instruct-GPTQ-Int4", trust_remote_code=True)
# print(tokenizer)
# use bf16
# model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B-Chat", device_map="auto", trust_remote_code=True, bf16=True).eval()
# use fp16
# model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B-Chat", device_map="auto", trust_remote_code=True, fp16=True).eval()
# use cpu only
# model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B-Chat", device_map="cpu", trust_remote_code=True).eval()
# use auto mode, automatically select precision based on the device.
# model = AutoModelForCausalLM.from_pretrained("hf-models/Qwen2-72B-Instruct-GPTQ-Int4", device_map="auto", trust_remote_code=True).eval()
# print(model)

# Specify hyperparameters for generation. But if you use transformers>=4.32.0, there is no need to do this.
# model.generation_config = GenerationConfig.from_pretrained("Qwen/Qwen-7B-Chat", trust_remote_code=True) # 可指定不同的生成长度、top_p等相关超参

# 第一轮对话 1st dialogue turn
# response, history = model.chat(tokenizer, "你好", history=None)
# print(response)
# 你好！很高兴为你提供帮助。

# 第二轮对话 2nd dialogue turn
# response, history = model.chat(tokenizer, "给我讲一个年轻人奋斗创业最终取得成功的故事。", history=history)
# print(response)
# 这是一个关于一个年轻人奋斗创业最终取得成功的故事。
# 故事的主人公叫李明，他来自一个普通的家庭，父母都是普通的工人。从小，李明就立下了一个目标：要成为一名成功的企业家。
# 为了实现这个目标，李明勤奋学习，考上了大学。在大学期间，他积极参加各种创业比赛，获得了不少奖项。他还利用课余时间去实习，积累了宝贵的经验。
# 毕业后，李明决定开始自己的创业之路。他开始寻找投资机会，但多次都被拒绝了。然而，他并没有放弃。他继续努力，不断改进自己的创业计划，并寻找新的投资机会。
# 最终，李明成功地获得了一笔投资，开始了自己的创业之路。他成立了一家科技公司，专注于开发新型软件。在他的领导下，公司迅速发展起来，成为了一家成功的科技企业。
# 李明的成功并不是偶然的。他勤奋、坚韧、勇于冒险，不断学习和改进自己。他的成功也证明了，只要努力奋斗，任何人都有可能取得成功。

# 第三轮对话 3rd dialogue turn
# response, history = model.chat(tokenizer, "给这个故事起一个标题", history=history)
# print(response)
# 《奋斗创业：一个年轻人的成功之路》

# def chat(message, history=None):
#     response, history = model.chat(tokenizer, message, history=history)
#     return response, history



command = ["python", "openai_api_server.py"]
# command = ["python", "--version"]
process = subprocess.Popen(
    command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

print("启动进程 openai_api_server.py ")

def read_process_output(process):
    """Reads the process output and prints it."""
    while True:
        print("读取进程输出  --- ")
        output = process.stdout.readline()
        if output == '' and process.poll() is not None:
            break
        if output:
            print("API Server: "+output.strip())



def read_process_error_output(process):
    """Reads the process output and prints it."""
    while True:
        print("读取进程错误  --- ")
        output = process.stderr.readline()
        if output == '' and process.poll() is not None:
            break
        if output:
            print("API Server error: "+output.strip())


# 启动一个线程读取 subprocess 输出
thread = threading.Thread(target=read_process_output, args=(process,))
print("启动线程读取 subprocess 输出 ")
thread.start()

# 启动一个线程读取 subprocess 输出
thread_error = threading.Thread(target=read_process_error_output, args=(process,))
print("启动线程读取 subprocess 错误输出 ")
thread_error.start()


system_prompt = {
    "role": "system",
    "content":
    "你是一个充满想象力的作家, 总是能在很短的篇幅中描绘生动形象的故事。" +
    "请根据夏商周的历史背景、神话传说和考古发现，创作一个关于一位古代英雄的冒险故事。" + 
    "故事背景设定在夏商时期。在一个对人类社会造成几乎毁灭性打击的大洪水过后，幸存下来的华夏民族开始重建并壮大自己的文明。" + 
    "故事跨域百年的时间，从夏朝开始，到商朝结束。" + 
    "故事的主题是人类不同族群的冲突与融合展现的文明的发展。" + 
    "主角名字是奇乙, 是寺人家族的年轻一代，和哥哥新甲对于在大洪水后的世界如何建立新秩序有不同的看法。" + 
    "他们常见因为不同见解而发生争执，于是有一天他们各自收到了禹王的命令，让哥哥新甲往东边探索， 让弟弟奇乙往西边探索。" + 
    "最终他们一起努力，为重建文明和实现民族融合做出贡献。" + 
    "这个故事带着读者跟随弟弟奇乙的探索脚本，踏上了充满未知和危险的旅程。在旅途中，主人公将面临各种挑战，包括敌对势力的阻挠、自然环境的考验以及内心的挣扎。" + 
    "通过结合历史资料、神话传说和考古发现，为故事增添真实感和神秘感。" + 
    "请确保故事情节紧凑，角色刻画丰富，并融入一些幽默元素，以增加故事的趣味性。" + 
    "最终，主人公能否成功找找到文明的生存之道，掌握自然界运行的规律，传播华夏文明，实现种族融合，以及他在旅途中所经历的冒险和成长，将成为故事的核心主题。" + 
    "你需要根据当前的故事状态来生成故事。"
}

import openai_api_request

def chat(message, history):
    messages = [system_prompt]
    for msg in history:
        messages.append({"role": "user", "content": str(msg[0])})
        messages.append({"role": "assistant", "content": str(msg[1])})

    messages.append({"role": "user", "content": str(message)})

    complete_message = ''

    res = openai_api_request.simple_chat(messages=messages, use_stream=False)
    print("res: ")
    print(res)

    delta_content = res.choices[0].message.content
    complete_message += delta_content
    
    # for chunk in res:
    #     print("chunk: ")
    #     print(chunk)
    #     delta_content = chunk.choices[0].delta.content
    #     complete_message += delta_content
        # print(delta_content, end='')  # 不换行拼接输出当前块的内容
        #  yield complete_message  # gradio 需要返回完整可迭代内容
    print(message)
    print("\nComplete message:", complete_message)
    return complete_message
